import os.path
import shutil
import string
import sys
import time

from multiprocessing import Pool, cpu_count


printable = string.ascii_letters + string.digits + string.punctuation


def RTxt(path):
  with open(path, encoding='utf8') as f:
    return f.read()


def WTxt(path, content):
  with open(path, 'w', encoding='utf8') as f:
    f.write(content)


def ATxt(path, content):
  with open(path, 'a', encoding='utf8') as f:
    f.write(content)


def RTxtLines(path):
  lines = RTxt(path).split('\n')
  return lines


class PlaceholderEnum:
  UPPER_CASE = r'\U'
  LOWER_CASE = r'\L'
  NUMBER = r'\N'
  ALL_CHAR = r'\A'
  SPECIAL_CHAR = r'\S'
  DICTIONARY = r'\D'


def GetPasswdSeq(passwdTemplate):

  def ZhuanYiCollect(seq, append=False):
      placeholderC = '\\' + c
      if placeholderC == PlaceholderEnum.NUMBER:
        if append:
          seq.append(list(string.digits))
        else:
          seq += list(string.digits)
      elif placeholderC == PlaceholderEnum.LOWER_CASE:
        if append:
          seq.append(list(string.ascii_lowercase))
        else:
          seq += list(string.ascii_lowercase)
      elif placeholderC == PlaceholderEnum.UPPER_CASE:
        if append:
          seq.append(list(string.ascii_uppercase))
        else:
          seq += list(string.ascii_uppercase)
      elif placeholderC == PlaceholderEnum.ALL_CHAR:
        if append:
          seq.append(list(printable))
        else:
          seq += list(printable)
      elif placeholderC == PlaceholderEnum.SPECIAL_CHAR:
        if append:
          seq.append(list(string.punctuation))
        else:
          seq += list(string.punctuation)
      elif placeholderC == PlaceholderEnum.DICTIONARY:
        # 字典路径
        if deepKwargs.get('dictionary', False):
          if append:
            seq.append(RTxtLines(''.join(deepKwargs['dictionarySeq'])))
          else:
            seq += RTxtLines(''.join(deepKwargs['dictionarySeq']))
          deepKwargs['dictionary'] = False
        else:
          deepKwargs['dictionary'] = True
          deepKwargs['dictionarySeq'] = []
      else:
        seq.append(c)

  zhuanYi = False
  chi = False
  deepKwargs = {}
  chiSeq = []
  passwdSeq = []
  for c in passwdTemplate:
    if zhuanYi:
      # 当前字符是转义字符
      if chi:
        # 收集池
        ZhuanYiCollect(chiSeq)
      else:
        ZhuanYiCollect(passwdSeq, append=True)
      zhuanYi = False
    else:
      if c == '\\':
        # 转义
        zhuanYi = True
      elif c == '[' and not chi:
        # 池开始
        chi = True
        chiSeq = []
      elif c == ']' and chi:
        # 池结束
        chi = False
        passwdSeq.append(chiSeq)
      else:
        if deepKwargs.get('dictionary', False):
          deepKwargs['dictionarySeq'].append(c)
        else:
          if chi:
            chiSeq.append(c)
          else:
            passwdSeq.append(c)
  return passwdSeq


def GetPasswords(args):
  passwdSeq, frontSeq, deep, index, minLength = args
  if frontSeq is None:
    frontSeq = []
  result = []
  for elementId in range(len(passwdSeq)):
    element = passwdSeq[elementId]
    if type(element) == list:
      if deep > 0:
        for e in element:
          result += GetPasswords(
            (passwdSeq[elementId + 1:], frontSeq + [e], deep + 1, '{}{}'.format(index, element.index(e)), minLength)
          )
        if len(result) > 5000000 or deep == 1:
          WTxt('temp/key_{}.tmp'.format(index), '\n'.join(result))
          return []
        return result
      else:
        argsList = []
        for e in element:
          argsList.append(
            (passwdSeq[elementId+1:], frontSeq + [e], deep + 1, '{}{}'.format(index, element.index(e)), minLength)
          )
        with Pool(processes=min(len(argsList), taskCount)) as pool:
          resultList = pool.map(GetPasswords, argsList)
        for r in resultList:
          result += r
        if result:
          WTxt('temp/key_root.tmp', '\n'.join(result))
      return []
    frontSeq.append(element)
  password = ''.join(frontSeq)
  if minLength > 0 and len(password) < minLength:
    return []
  result.append(password)
  return result


def GetArgs(key, default=None):
  if key not in sys.argv:
    return default
  vIndex = sys.argv.index(key)
  if vIndex + 1 < len(sys.argv):
    return sys.argv[vIndex + 1]
  return default


def GetPasswdNumber(passwdSeq):
  # print('passwdSeq=', passwdSeq)
  result = 1
  for elementId in range(len(passwdSeq)):
    element = passwdSeq[elementId]
    if type(element) == list:
      result *= len(element)
  return result


def Merge():
  files = os.listdir('temp/')
  for file in files[:-1]:
    print('\r合并进度: {} / {}'.format(files.index(file), len(files)), end='')
    ATxt(savePath, RTxt('temp/{}'.format(file)) + '\n')
  print('\r合并进度: {} / {}'.format(len(files), len(files)), end='')
  ATxt(savePath, RTxt('temp/{}'.format(files[-1])))
  shutil.rmtree('temp')


class Creator:
  def __init__(self, passwdTemplate):
    self.passwdTemplate = passwdTemplate

  def Working(self):
    startTime = time.time()
    passwdSeq = GetPasswdSeq(self.passwdTemplate)
    print('预计密码数量：{}'.format(GetPasswdNumber(passwdSeq)))
    GetPasswords((passwdSeq, None, 0, '', passwdMinLength))
    print('开始合并缓存文件！')
    Merge()
    print('\n字典文件生成完成！用时：{} 秒'.format(round(time.time() - startTime), 2))


if __name__ == '__main__':
  if '-h' in sys.argv:
    print(r'''
    python3 main.py [-h] [-e 字典生成表达式] [-o 字典输出路径] [-l 密码最短长度]
    -e: 字典生成表达式
      \U: 大写字母掩码
      \L: 小写字母掩码
      \N: 数字掩码
      \S: 特殊字符掩码
      \A: 所有字符掩码
      \D字典路径\D: 读取其它字典文件内密码
      [字符或掩码]: 指定字符组合的掩码
      例如:
      python3 main.py -e 123[\U\L]\NAbc 
      python3 main.py -e 123[\N\S]\Dpasswds.txt\D\N 
      
    -l: 根据表达式生成的密码最小长度（低于长度的密码将被丢弃），大于0时生效，默认为-1
      例如:
      python3 main.py -e 123[\N\S]\Dpasswds.txt\D\N -l 12
      
    -o: 根据表达式生成的密码保存路径
      例如:
      python3 main.py -e 123[\U\L]\NAbc -o output.txt
      
    -p: 设置最大线程
      
    -h: 显示帮助信息
    ''')
    exit(0)
  savePath = GetArgs('-o', 'output.txt')
  if os.path.exists(savePath):
    os.remove(savePath)
  if not os.path.exists('temp'):
    os.mkdir('temp')
  expression = GetArgs('-e', None)
  passwdMinLength = int(GetArgs('-l', '-1'))
  taskCount = int(GetArgs('-p', '-1'))
  if taskCount == -1:
    taskCount = cpu_count()
  if expression is None:
    print('''
    未用 -e 指定表达式
    更多信息请使用 python3 main.py -h 参考帮助信息
    ''')
    exit(0)
  creator = Creator(expression)
  creator.Working()
